'''
yolov5 aidlux包装类
'''
import os
import cv2
import numpy as np
import aidlite_gpu

from utils import preprocess_img, detect_postprocess, convert_shape

class Yolov5(object):
    coco_class = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
        'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
        'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
        'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
        'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
        'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
        'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
        'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
        'hair drier', 'toothbrush']

    def __init__(self, model_path:str, input_shape:tuple, output_shape:tuple, device:int, int8:bool, 
                        input_scale:float or list=None, input_zero_point:float or list=None,
                        output_sacle:float or list=None, output_zero_point:float or list=None,
                        conf_thres:float=0.5, iou_thres:float=0.4, interested_classid:list=None) -> None:
        """
        Yolov5 aidlux包装类
        如果是int8模型,输入图像会做 img / input_scale + input_zero_point 处理, 输出会做 (pred - output_zero_point)*output_sacle
        :param model_path: 模型路径
        :param input_shape: 模型输入shape, eg: (1, 320, 320, 3)
        :param output_shape: 模型输出shape, eg: (1, 8300, 85)
        :param device: 运行设备, -1:cpu, 0:gpu, 1:cpu+gpu, 2:dsp
        :param int8: 是否是int8模型
        :param input_scale: 输入scale
        :param input_zero_point: 输入zero_point
        :param output_sacle: 输出scale
        :param output_zero_point: 输出zero_point
        :param conf_thres: 置信度阈值
        :param iou_thres: nms iou阈值
        :param interested_classid: 输出类别的id,应为list,为None时则输出所有
        :return None
        """
        assert os.path.exists(model_path), f"{model_path} 不存在"
        if interested_classid is not None:
            assert isinstance(interested_classid, list), f"interested_classid 应为list或None"
        self.model = self.bulid_model(model_path, input_shape, output_shape, device, int8)
        self.input_shape = input_shape
        self.output_shape = output_shape
        self.interested_classid = interested_classid
        self.int8 =int8
        self.input_scale = input_scale
        self.input_zero_point = input_zero_point
        self.output_sacle = output_sacle
        self.output_zero_point = output_zero_point
        self.conf_thres = conf_thres
        self.iou_thres = iou_thres

    def bulid_model(self, model_path:str, input_shape:tuple, output_shape:tuple, device=-1, int8=False):
        '''
        构建在aidlux上运行的模型
        :param model_path: 模型路径
        :param input_shape: 模型输入shape, eg: (1, 320, 320, 3)
        :param output_shape: 模型输出shape, eg: (1, 8300, 85), [(1, 520, 520, 85), (1, 32, 32, 85)]
        :param device: 运行设备, -1:cpu, 0:gpu, 1:cpu+gpu, 2:dsp
        :param int8: 是否是int8模型
        :return aidlte_model
        '''
        aidlite = aidlite_gpu.aidlite(0)
        aidlite.FAST_ANNModel(model_path, 
                                convert_shape(input_shape, int8), 
                                convert_shape(output_shape, int8), 
                                4, 
                                device)
        return aidlite

    def invoke(self, ori_img:np.ndarray) -> np.ndarray:
        '''
        推理函数
        :param ori_img: 预测图像
        :return: 预测坐标框信息, np.ndarray(N, 6), xywh、conf、cls_id
        '''
        det_pred = None
        if self.int8:
            input_img = cv2.resize(ori_img, (self.input_shape[1], self.input_shape[2]))
            if self.input_scale is not None or self.input_zero_point is not None:
                input_img = input_img/255
            if self.input_scale is not None:
                input_img = input_img/self.input_scale
            if self.input_zero_point is not None:
                input_img = input_img + self.input_zero_point
            self.model.setInput_Int8(input_img)
            self.model.invoke()
            det_pred = self.model.getOutput_Int8().reshape(*self.output_shape)[0].astype(np.float32)

            if self.output_zero_point is not None:
                det_pred = det_pred - self.output_zero_point
            if self.output_sacle is not None:
                det_pred = det_pred * self.output_sacle
        else:
            input_img = preprocess_img(ori_img, 
                                        target_shape=(self.input_shape[1], self.input_shape[2]),
                                        div_num=255,
                                        means=None,
                                        stds=None)
            self.model.setInput_Float32(input_img)
            self.model.invoke()
            det_pred = self.model.getOutput_Float32().reshape(*self.output_shape)[0]

        det_pred = detect_postprocess(det_pred, ori_img.shape, [self.input_shape[1], self.input_shape[2], 3], 
                                        conf_thres=0.5, iou_thres=0.4, interested_classid=self.interested_classid)
        return det_pred

    def draw_pic(self, img:np.ndarray, det_pred:np.ndarray) -> np.ndarray:
        '''
        绘图函数
        :param img: 图像
        :param det_pred: 预测结果 np.ndarray(N, 6), 坐标框信息, x1y1x2y2、conf、cls_id
        :return: 完成预测结果绘制的图像
        '''
        img = img.astype(np.uint8)
        color_step = int(255/len(det_pred))
        for i in range(len(det_pred)):
            x1, y1, x2, y2, conf, cls_id  = [int(t) for t in det_pred[i]]
            cv2.putText(img, f'{self.coco_class[cls_id]}', (x1, y1-6), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
            cv2.rectangle(img, (x1, y1), (x2, y2), (0, cls_id*color_step, 255-cls_id*color_step),thickness = 2)
        return img